#include <string.h>
#include <stdio.h>
#include "rijndael.h"

typedef uint8_t sm_key_t[16];
typedef uint8_t bd_addr_t[6];

void dump(const char* name, const unsigned char* data, int len){
    printf("%s: ", name);
    for (int i = 0; i < len; i++){
        printf("%02x ", data[i]);
    }
    printf("\n");
}

void reverse_bytes(const uint8_t * src, uint8_t * dest, int len){
    int i;
    for (i = 0; i < len; i++)
        dest[len - 1 - i] = src[i];
}
void reverse_56(const uint8_t * src, uint8_t * dest){
    reverse_bytes(src, dest, 7);
}
void reverse_128(const uint8_t * src, uint8_t * dest){
    reverse_bytes(src, dest, 16);
}

void aes128_calc(const uint8_t * key, const uint8_t * plaintext, uint8_t * ciphertext){
    uint32_t rk[RKLENGTH(KEYBITS)];
    int nrounds = rijndaelSetupEncrypt(rk, &key[0], KEYBITS);
    rijndaelEncrypt(rk, nrounds, plaintext, ciphertext);
}
static void sm_c1_t1(sm_key_t r, uint8_t preq[7], uint8_t pres[7], uint8_t iat, uint8_t rat, uint8_t * t1){

    // p1 = pres || preq || rat’ || iat’
    // "The octet of iat’ becomes the least significant octet of p1 and the most signifi-
    // cant octet of pres becomes the most significant octet of p1.
    // For example, if the 8-bit iat’ is 0x01, the 8-bit rat’ is 0x00, the 56-bit preq
    // is 0x07071000000101 and the 56 bit pres is 0x05000800000302 then
    // p1 is 0x05000800000302070710000001010001."

    sm_key_t p1;
    reverse_56(pres, &p1[0]);
    reverse_56(preq, &p1[7]);
    p1[14] = rat;
    p1[15] = iat;

    // t1 = r xor p1
    int i;
    for (i=0;i<16;i++){
        t1[i] = r[i] ^ p1[i];
    }
}
static void sm_c1_t3(sm_key_t t2, bd_addr_t ia, bd_addr_t ra, uint8_t * t3){
    // p2 = padding || ia || ra
   // "The least significant octet of ra becomes the least significant octet of p2 and
   // the most significant octet of padding becomes the most significant octet of p2.
   // For example, if 48-bit ia is 0xA1A2A3A4A5A6 and the 48-bit ra is
   // 0xB1B2B3B4B5B6 then p2 is 0x00000000A1A2A3A4A5A6B1B2B3B4B5B6.

   sm_key_t p2;
   // cppcheck-suppress uninitvar ; p2 is reported as uninitialized
   memset(p2, 0, 16);
   (void)memcpy(&p2[4], ia, 6);
   (void)memcpy(&p2[10], ra, 6);

   // c1 = e(k, t2_xor_p2)
   int i;
   for (i=0;i<16;i++){
       t3[i] = t2[i] ^ p2[i];
   }
}

static void sm_s1_r_prime(sm_key_t r1, sm_key_t r2, uint8_t * r_prime){
    (void)memcpy(&r_prime[8], &r2[8], 8);
    (void)memcpy(&r_prime[0], &r1[8], 8);
}

bd_addr_t local_addr = {0xe2, 0x00, 0x20, 0x00, 0x13, 0x9e};
bd_addr_t remote_addr = {0xff, 0xff, 0x56, 0x39, 0xb6, 0xa9};
uint8_t local_addr_type = 0x00;
uint8_t remote_addr_type = 0x01;

//-----------------------------------copy from hci start-----------------------------------
unsigned char pair_req[] = {0x1, 0x1, 0x0, 0x1, 0x10, 0x5, 0x5};
unsigned char pair_rsp[] = {0x2, 0x3, 0x0, 0x1, 0x10, 0x1, 0x1};

#if 0
unsigned char plain1[] = {0xf7, 0x58, 0x7e, 0xfc, 0x18, 0x2b, 0x2b, 0xe3, 0xc7, 0x3a, 0xa, 0xea, 0x8f, 0xaf, 0xbf, 0xa5};
unsigned char enc_result1[] = {0x6f, 0xe8, 0xad, 0x13, 0x1a, 0x98, 0xed, 0x42, 0xca, 0x97, 0xfa, 0x95, 0x92, 0x8d, 0xcb, 0x3f};
unsigned char plain2[] = {0xc6, 0x5e, 0x94, 0x45, 0xe5, 0x67, 0x73, 0x51, 0xca, 0xb7, 0xfa, 0x77, 0x92, 0x8d, 0xcb, 0x3f};
unsigned char enc_result2[] = {0x6b, 0xf5, 0x3d, 0x52, 0xa4, 0x36, 0x6f, 0x30, 0x3d, 0x83, 0xf, 0x7f, 0x5c, 0x6e, 0x72, 0x7};
unsigned char send_confirm[] = {0x6b, 0xf5, 0x3d, 0x52, 0xa4, 0x36, 0x6f, 0x30, 0x3d, 0x83, 0xf, 0x7f, 0x5c, 0x6e, 0x72, 0x7};
unsigned char send_random[] = {0xf7, 0x59, 0x7f, 0xfd, 0x18, 0x2a, 0x3b, 0xe6, 0xc2, 0x38, 0x9, 0xea, 0x8e, 0xbf, 0xbe, 0xa4};

unsigned char recv_confirm[] = {0x5b, 0x18, 0x63, 0xb2, 0xe9, 0xb2, 0xfd, 0xc1, 0xc8, 0x21, 0x43, 0xc5, 0xc, 0x13, 0x57, 0x99};
unsigned char recv_random[] = {0x4c, 0xb3, 0x2e, 0xc7, 0xe5, 0x3, 0xcd, 0x3d, 0x79, 0xe9, 0x5f, 0xac, 0x5b, 0x43, 0x5d, 0xec};
unsigned char remote_plain1[] = {0x4c, 0xb2, 0x2f, 0xc6, 0xe5, 0x2, 0xdd, 0x38, 0x7c, 0xeb, 0x5c, 0xac, 0x5a, 0x53, 0x5c, 0xed};
unsigned char remote_enc1[] = {0x7a, 0xc8, 0x4d, 0xbb, 0x21, 0xf7, 0xc3, 0x6d, 0xb9, 0x19, 0x38, 0x97, 0x69, 0xd, 0x53, 0x62};
unsigned char remote_plain2[] = {0x8c, 0xd1, 0xb1, 0xb5, 0xde, 0x8, 0x5d, 0x7e, 0xb9, 0x39, 0x38, 0x75, 0x69, 0xd, 0x53, 0x62};
unsigned char remote_enc2[] = {0xaf, 0xa1, 0x88, 0x2b, 0x4a, 0xb7, 0x66, 0x15, 0x5d, 0x1e, 0x61, 0x6d, 0x78, 0xc, 0xe1, 0xa2};

unsigned char stk_plain[] = {0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0};
unsigned char stk[] = {0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0};
#else
unsigned char plain1[] = {0xa7, 0x16, 0xc, 0x75, 0x54, 0xe3, 0xb7, 0x10, 0xa9, 0x30, 0xa6, 0x90, 0xf7, 0x93, 0xd, 0xf9};
unsigned char enc_result1[] = {0xce, 0xf1, 0x77, 0x6f, 0xc7, 0x99, 0x9d, 0x15, 0x53, 0x6e, 0x4a, 0xe0, 0xe6, 0x79, 0xd8, 0xc3};
unsigned char plain2[] = {0x67, 0x47, 0x4e, 0x39, 0x38, 0x66, 0x3, 0x6, 0x53, 0x4e, 0x4a, 0x2, 0xe6, 0x79, 0xd8, 0xc3};
unsigned char enc_result2[] = {0xb5, 0x2b, 0xb6, 0x41, 0xe5, 0x4e, 0x54, 0x30, 0x35, 0xc7, 0xd8, 0x61, 0x40, 0x59, 0xaa, 0xe3};
unsigned char send_confirm[] = {0xb5, 0x2b, 0xb6, 0x41, 0xe5, 0x4e, 0x54, 0x30, 0x35, 0xc7, 0xd8, 0x61, 0x40, 0x59, 0xaa, 0xe3};
unsigned char send_random[] = {0xa7, 0x17, 0xd, 0x74, 0x54, 0xe2, 0xa7, 0x15, 0xac, 0x32, 0xa5, 0x90, 0xf6, 0x83, 0xc, 0xf8};

unsigned char recv_confirm[] = {0xb9, 0x7b, 0x13, 0x1e, 0x1a, 0x83, 0xf6, 0x82, 0xd2, 0x75, 0x65, 0xf3, 0x96, 0xbe, 0xe1, 0xa9};
unsigned char recv_random[] = {0x74, 0x52, 0x73, 0x3d, 0xa9, 0x6c, 0x50, 0x42, 0x63, 0xa6, 0x90, 0xee, 0x0, 0x15, 0x76, 0x21};
unsigned char remote_plain1[] = {0x74, 0x53, 0x72, 0x3c, 0xa9, 0x6d, 0x40, 0x47, 0x66, 0xa4, 0x93, 0xee, 0x1, 0x5, 0x77, 0x20};
unsigned char remote_enc1[] = {0x39, 0x1d, 0x1d, 0xd9, 0xf8, 0xcc, 0x3d, 0xab, 0x4f, 0x63, 0xe6, 0xd9, 0x8d, 0x3a, 0x7f, 0x10};
unsigned char remote_plain2[] = {0x90, 0xab, 0x24, 0x8f, 0x7, 0x33, 0xa3, 0xb8, 0x4f, 0x43, 0xe6, 0x3b, 0x8d, 0x3a, 0x7f, 0x10};
unsigned char remote_enc2[] = {0xb9, 0x7b, 0x13, 0x1e, 0x1a, 0x83, 0xf6, 0x82, 0xd2, 0x75, 0x65, 0xf3, 0x96, 0xbe, 0xe1, 0xa9};

unsigned char stk_plain[] = {0xa7, 0x17, 0xd, 0x74, 0x54, 0xe2, 0xa7, 0x15, 0x74, 0x52, 0x73, 0x3d, 0xa9, 0x6c, 0x50, 0x42};
unsigned char stk[] = {0x0, 0x3c, 0x97, 0xe1, 0x2f, 0xb0, 0x13, 0xec, 0xbf, 0x91, 0x80, 0x59, 0x99, 0x45, 0x52, 0xf1};
#endif
//------------------------------------copy from hci end------------------------------

void le_encrypt(const unsigned char* plain, unsigned char* out){
    unsigned char key[16] = {0};
    unsigned char plain_rev[16] = {0};
    unsigned char out_rev[16] = {0};

    reverse_128(plain, plain_rev);

    aes128_calc(key, plain_rev, out_rev);

    reverse_128(out_rev, out);
}

void get_t1(unsigned char* key, unsigned char* t1){
    sm_key_t r;
    sm_key_t t1_rev;
    reverse_128(key, r);

    sm_c1_t1(r, pair_req, pair_rsp, local_addr_type, remote_addr_type, t1_rev);

    reverse_128(t1_rev, t1);
}

void get_stk(unsigned char* master_random, unsigned char* slave_random, unsigned char* stk){
    sm_key_t r1;
    sm_key_t r2;
    sm_key_t r_prime;
    sm_key_t prime_rev;

    reverse_128(master_random, r1);
    reverse_128(slave_random, r2);

    sm_s1_r_prime(r2, r1, prime_rev);
    reverse_128(prime_rev, r_prime);

    dump("r_prime", r_prime, 16);
    le_encrypt(r_prime, stk);
}

void get_t3(unsigned char* key, unsigned char* t3){
    sm_key_t r;
    sm_key_t t3_rev;
    reverse_128(key, r);


    sm_c1_t3(r,local_addr, remote_addr, t3_rev);

    reverse_128(t3_rev, t3);
}

int main(){
    unsigned char out[16] = {0};
    unsigned char t3[16] = {0};
    unsigned char t1[16] = {0};

    dump("local random", send_random, 16);

    get_t1(send_random, t1);
    dump("local t1", t1, 16);
    if(!memcmp(t1, plain1, 16)){
        printf("local t1 ok\n");
    }else{
        printf("local t1 fail\n");
    }

    le_encrypt(t1, out);
    dump("enc local t1", out, 16);

    get_t3(out, t3);
    dump("local t3", t3, 16);
    if(!memcmp(t3, plain2, 16)){
        printf("local t3 ok\n");
    }else{
        printf("local t3 fail\n");
    }

    le_encrypt( t3, out);
    dump("local confirm", out, 16);

    if(!memcmp(out, send_confirm, 16)){
        printf("local confirm ok\n");
    }else{
        printf("local confirm fail\n");
    }

    printf("----remote----\n");
    dump("remote random", recv_random, 16);

    get_t1(recv_random, t1);
    dump("remote t1", t1, 16);
    if(!memcmp(t1, remote_plain1, 16)){
        printf("remote t1 ok\n");
    }else{
        printf("remote t1 fail\n");
    }

    le_encrypt(t1, out);
    dump("enc remote t1", out, 16);

    get_t3(out, t3);
    dump("remote t3", t3, 16);
    if(!memcmp(t3, remote_plain2, 16)){
        printf("remote t3 ok\n");
    }else{
        printf("remote t3 fail\n");
    }

    le_encrypt(t3, out);
    dump("calc remote confirm", out, 16);

    if(!memcmp(out, recv_confirm, 16)){
        printf("remote confirm ok\n");
    }else{
        printf("remote confirm fail\n");
    }

    printf("----stk----\n");
    get_stk(send_random, recv_random, out);
    dump("stk", out, 16);
    if(!memcmp(out, stk, 16)){
        printf("stk ok\n");
    }else{
        printf("stk fail\n");
    }

    printf("----xor t1----\n");
    for (int i = 0; i < 16; i++){
        out[i] = send_random[i] ^ plain1[i];
    }
    dump("local  xor", out, 16);

    for (int i = 0; i < 16; i++){
        out[i] = recv_random[i] ^ remote_plain1[i];
    }
    dump("remote xor", out, 16);

    printf("----xor t3----\n");
    for (int i = 0; i < 16; i++){
        out[i] = enc_result1[i] ^ plain2[i];
    }
    dump("local  xor", out, 16);

    for (int i = 0; i < 16; i++){
        out[i] = remote_enc1[i] ^ remote_plain2[i];
    }
    dump("remote xor", out, 16);
}